package com.means.polymerize.assembly.filter;

import cn.hutool.core.io.IoUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.ReflectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HTMLFilter;
import lombok.SneakyThrows;
import org.springframework.http.MediaType;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.ServletRequest;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.Map;
import java.util.stream.IntStream;

/**
 * Xss 请求 Wrapper
 *
 * @author xuhaifneg
 * @date 2021.05.05
 */
public class XssRequestWrapper extends HttpServletRequestWrapper {
    private static final String ENCODE_QUOTE = "encodeQuotes";
    /**
     * 基于线程级别的 HTMLFilter 对象，因为它线程非安全
     */
    private static final ThreadLocal<HTMLFilter> HTML_FILTER = ThreadLocal.withInitial(() -> {
        final HTMLFilter htmlFilter = new HTMLFilter();
        /* 反射修改 encodeQuotes 属性为 false，避免 " 被转移成 &quot; 字符 */
        ReflectUtil.setFieldValue(htmlFilter, ENCODE_QUOTE, false);
        return htmlFilter;
    });

    /**
     * Instantiates a new Xss request wrapper.
     *
     * @param request the request
     */
    public XssRequestWrapper(HttpServletRequest request) {
        super(request);
    }

    private static String filterXss(String content) {
        if (StrUtil.isEmpty(content)) {
            return content;
        }
        return HTML_FILTER.get().filter(content);
    }

    // ========== IO 流相关 ==========

    private static boolean isJsonRequest(ServletRequest request) {
        return StrUtil.startWithIgnoreCase(request.getContentType(), MediaType.APPLICATION_JSON_VALUE);
    }

    @Override
    public BufferedReader getReader() {
        return IoUtil.toBuffered(IoUtil.getReader(this.getInputStream(), CharsetUtil.CHARSET_UTF_8));
    }

    // ========== Param 相关 ==========

    @Override
    @SneakyThrows(value = {IOException.class})
    public ServletInputStream getInputStream() {
        /* 如果非 json 请求，不进行 Xss 处理 */
        if (!isJsonRequest(this)) {
            return super.getInputStream();
        }
        /* 读取内容，并过滤 */
        final String content = filterXss(IoUtil.readUtf8(super.getInputStream()));
        final ByteArrayInputStream newInputStream = IoUtil.toStream(content.getBytes());
        return new ServletInputStream() {
            @Override
            public int read() {
                return newInputStream.read();
            }

            @Override
            public boolean isFinished() {
                return true;
            }

            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public void setReadListener(ReadListener readListener) {
            }
        };
    }

    @Override
    public String getParameter(String name) {
        return filterXss(super.getParameter(name));
    }

    @Override
    public String[] getParameterValues(String name) {
        final String[] values = super.getParameterValues(name);
        if (ArrayUtil.isEmpty(values)) {
            return values;
        }
        IntStream.range(0, values.length).forEach(i -> values[i] = filterXss(values[i]));
        return values;
    }

    // ========== Header 相关 ==========

    @Override
    public Map<String, String[]> getParameterMap() {
        final Map<String, String[]> valueMap = super.getParameterMap();
        if (MapUtil.isEmpty(valueMap)) {
            return valueMap;
        }
        valueMap.forEach((key, values) ->
                IntStream.range(0, values.length).forEach(i -> values[i] = filterXss(values[i])));
        return valueMap;
    }

    @Override
    public String getHeader(String name) {
        return filterXss(super.getHeader(name));
    }
}
